#!/usr/bin/python

# after ase line add p-val of chi-square test , minor allele frequency, and decode (1,2,3,0)
# 0: low power (not enough read depth), 1: homoallele (allele 1), 2: homoallele (allele 2), 3: heteroallele (allele 1 = allele 2)
# usage: python ase_eval_alpha_depth.py AN00493_ba9_selected_snp_10.ase 0.05 10 > AN00493_ba9_selected_snp_10.ase-decode_0.05_10
# At mac, ase_eval_alpha_depth.py AN00493_ba9_selected_snp_10.ase 0.05 10 > AN00493_ba9_selected_snp_10.ase-decode cause error to load module


# install numpy and scipy from http://www1.i2r.a-star.edu.sg/~lins/codes/python.html
# reference sites for chi-square test
# http://stackoverflow.com/questions/9330114/chi-squared-test-in-python
# http://stackoverflow.com/questions/13913572/difference-between-scipy-stats-mstats-chisquare-and-scipy-stats-mstats-in-python

import sys

# load module for chi-square test
# to use python package, type "module load python" before run this program
import numpy as np
import scipy.stats.mstats as mst

### input
file1 = sys.argv[1]
##file1 = "AN00493_ba9_selected_snp_10.ase"
data = open(file1, "r")

# alpha input
alpha_num = sys.argv[2]
alpha = float(alpha_num)

# read depth input
read_depth_num = sys.argv[3]
read_depth = float(read_depth_num)

for line in data:
    col = line.split()

##    # chi-square test
##    obs = np.array([int(col[5]), int(col[6])]) # observed values
##    exp = np.array([.5, .5]) * np.sum(obs)     # expected values
##    print mst.chisquare(obs, exp)

    # chi-square test -- simple way when exp is 1:1 ratio
    obs = np.array([float(col[5]), float(col[6])]) # observed values
    chi_p = mst.chisquare(obs)[1]
    
    # minor allele frequency
    if float(col[5]) + float(col[6]) == 0:   # Both alleles are not detectable
        feq = "NA"  
    else:
        feq = 100*min(float(col[5]), float(col[6]))/(float(col[5]) + float(col[6]))

    # decode  ## alpha is 0.01 for ASE
    if float(col[5]) + float(col[6]) < read_depth:    # low power (not enough read depth)
        decode = 0
    elif chi_p < alpha and float(col[5]) >= float(col[6]):  # homoallele (allele 1)
        decode = 1
    elif chi_p < alpha and float(col[5]) < float(col[6]):   # homoallele (allele 2)
        decode = 2
    elif chi_p >= alpha:  # heteroallele (allele 1 = allele 2)
        decode = 3
    print line.strip()+"\t"+str(chi_p)+"\t"+str(feq)+"\t"+str(decode)
    
data.close()
